/*
 * Copyright (c) 2012 Qunar.com. All Rights Reserved
 */

package com.pyb.common.utils;

import com.google.common.base.Strings;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;

import org.springframework.util.CollectionUtils;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;

/**
 * QpsControl
 *
 * @author chuchen.zhang
 * @date 16-10-8 下午1:51
 */
public class TrafficShaper {

    private static final ConcurrentMap<String, RateLimiter> resourceLimiterMap = Maps.newConcurrentMap();

    public static void updateResourceQps(String resource, double qps) {
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            limiter = RateLimiter.create(qps);
            RateLimiter putByOtherThread
                    = resourceLimiterMap.putIfAbsent(resource, limiter);
            if (putByOtherThread != null) {
                limiter = putByOtherThread;
            }
        }
        limiter.setRate(qps);
    }

    public static void removeResource(String resource) {
        resourceLimiterMap.remove(resource);
    }

    /**
     * 尝试qps控制
     *
     * @throws ResourceNotFoundException 无resource，抛出异常
     * @return 尝试成功/失败
     */
    public static boolean tryEnter(String resource) throws ResourceNotFoundException{
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            throw new ResourceNotFoundException(resource);
        }
        return limiter.tryAcquire();
    }

    /**
     * qps控制
     *
     * @throws RateLimitException 如果qps限制，抛出异常
     * @throws ResourceNotFoundException 无resource，抛出异常
     */
    public static void enterWithException(String resource) throws RateLimitException, ResourceNotFoundException {
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            throw new ResourceNotFoundException(resource);
        }
        if (!limiter.tryAcquire()) {
            throw new RateLimitException(resource);
        }
    }

    /**
     * 阻塞qps控制
     *
     * @throws ResourceNotFoundException 无resource，抛出异常
     */
    public static void blockingEnter(String resource) throws ResourceNotFoundException {
        RateLimiter limiter = resourceLimiterMap.get(resource);
        if (limiter == null) {
            throw new ResourceNotFoundException(resource);
        }
        limiter.acquire();
    }

    /**
     * 更新传入的map
     */
    public static void updateMany(Map<String, Double> updateInUseMap) {
        updateAndRemoveUnuse(updateInUseMap, null, null);
    }

    /**
     * 删除多个
     */
    public static void removeMany(String prefix, String suffix) {
        updateAndRemoveUnuse(null, prefix, suffix);
    }

    /**
     * 更新传入的map, 删除相关其它的rateLimitter；前缀、后缀都为empty时，不删除其它
     * @param updateInUseMap if empty, will release all with prefix & suffix
     * @param prefix 无前缀限定，传入null or ""
     * @param suffix 无后缀限定，传入null or ""
     */
    public static void updateAndRemoveUnuse(Map<String, Double> updateInUseMap, String prefix, String suffix) {
        boolean willUpdate = true;

        if (CollectionUtils.isEmpty(updateInUseMap)) {
            willUpdate = false;
        }

        boolean hasPrefix = !Strings.isNullOrEmpty(prefix);
        boolean hasSuffix = !Strings.isNullOrEmpty(suffix);
        boolean willRemove = hasPrefix || hasSuffix;

        if (willUpdate) {
            for (String resource : updateInUseMap.keySet()) {
                updateResourceQps(resource, updateInUseMap.get(resource));
            }
        }

        if (!willRemove) {
            return;
        }

        Set<String> resourceSetInCache = resourceLimiterMap.keySet();
        for (String resource : resourceSetInCache) {
            boolean isMatch = false;
            if (hasPrefix && hasSuffix) {
                if (resource.startsWith(prefix) && resource.endsWith(suffix)){
                    isMatch = true;
                }
            } else if ((hasPrefix && resource.startsWith(prefix)) || (hasSuffix && resource.endsWith(suffix))) {
                isMatch = true;
            }
            // match & not in update map, will remove
            if (isMatch && !updateInUseMap.containsKey(resource)) {
                removeResource(resource);
            }
        }
    }

    public static class RateLimitException extends Exception {

        private static final long serialVersionUID = 1L;

        private String resource;

        public String getResource() {
            return resource;
        }

        public RateLimitException(String resource) {
            super(resource + " should not be visited so frequently");
            this.resource = resource;
        }

        @Override
        public synchronized Throwable fillInStackTrace() {
            return this;
        }
    }

    public static class ResourceNotFoundException extends Exception {

        private static final long serialVersionUID = 1L;

        private String resource;

        public String getResource() {
            return resource;
        }

        public ResourceNotFoundException(String resource) {
            super(resource + " resource not found");
            this.resource = resource;
        }

        @Override
        public synchronized Throwable fillInStackTrace() {
            return this;
        }
    }
}
